"""
Frequency calculation module for periodic date operations.

This module provides base classes and implementations for calculating
periodic dates such as daily, weekly, monthly, quarterly, etc.
"""

from abc import ABC, abstractmethod
from functools import lru_cache

import pandas as pd
from typing import Union, List

from ..config import (
    DEFAULT_BEGIN_DATE,
    DEFAULT_END_DATE,
    DAILY_SEARCH_STEPS,
    WEEKLY_OFFSET,
    MONTHLY_OFFSET,
    QUARTERLY_OFFSET,
    YEARLY_OFFSET
)
from ..utils import dt_handle
from .trade_dt_handle import adjust_trade_dt, get_trade_dts, step_trade_dt


class FrequencyError(Exception):
    """Custom exception for frequency calculation errors."""
    pass


class RefreshBase(ABC):
    """
    获取调仓日期的基类.

    This abstract base class defines the interface for all frequency
    calculation classes including daily, weekly, monthly, quarterly, etc.

    Attributes:
        *args: Direction indicators, must be -1 or 1

    Examples:
        >>> class CustomFrequency(RefreshBase):
        ...     def get(self, begin_dt, end_dt):
        ...         # Custom logic here
        ...         return dates
    """

    def __init__(self, *args: int) -> None:
        """
        Initialize frequency with direction indicators.

        Args:
            *args: Direction indicators, must be -1 or 1

        Raises:
            ValueError: If args are not -1 or 1, or sum exceeds 1

        Examples:
            >>> freq = RefreshBase(1)  # Forward direction
            >>> freq = RefreshBase(-1)  # Backward direction
            >>> freq = RefreshBase(1, -1)  # Both directions
        """
        if not args:
            raise ValueError("At least one argument required")

        # Validate all arguments
        for arg in args:
            if arg not in (-1, 1):
                raise ValueError(f"Arguments must be -1 or 1, got {arg}")

        if abs(sum(args)) > 1:
            raise ValueError(f"Sum of arguments must not exceed 1, got {sum(args)}")

        self.args = sorted(args, reverse=True)

    @abstractmethod
    def get(self, begin_dt: Union[str, int], end_dt: Union[str, int]) -> pd.Series:
        """
        Get dates within the specified range.

        Args:
            begin_dt: Begin date (YYYYMMDD format or datetime)
            end_dt: End date (YYYYMMDD format or datetime)

        Returns:
            pandas Series of dates in YYYYMMDD format

        Raises:
            NotImplementedError: If not implemented by subclass
        """
        pass

    @lru_cache(maxsize=512)
    def next(self, date: Union[str, int], step: int = 1, adjust: bool = True) -> str:
        """
        获取下一个频率日期.

        Examples: Monthly(-1): 20180921 -> 20180928

        Args:
            date: Starting date (YYYYMMDD format or datetime)
            step: Number of steps forward (default: 1)
            adjust: Whether to adjust if date is exactly on a frequency boundary (default: True)

        Returns:
            Next date in the class frequency (YYYYMMDD format)

        Examples:
            >>> monthly = Monthly(1)
            >>> monthly.next('20240115')  # Get next month-end
            '20240131'

        Raises:
            FrequencyError: If calculation fails or step is invalid
        """
        if step < 1:
            raise FrequencyError(f"Step must be positive, got {step}")

        date = dt_handle(date)
        end_dt = step_trade_dt(date, DAILY_SEARCH_STEPS)
        df = self.get(date, end_dt).tolist()

        try:
            if df and df[0] == date:
                if adjust and len(df) > step:
                    return df[step]
                elif len(df) > step:
                    return df[step]

            if df and len(df) >= step:
                return df[step - 1]
            elif df:
                return df[-1]
            else:
                raise FrequencyError(f"No frequency dates found after {date}")
        except IndexError as e:
            if df:
                return df[-1]
            raise FrequencyError(f"Cannot get next date from {date}: {e}")

    @lru_cache(maxsize=512)
    def prev(self, date: Union[str, int], step: int = 1, adjust: bool = True) -> str:
        """
        获取上一个频率日期.

        Examples: Monthly(-1): 20180921 -> 20180831

        Args:
            date: Starting date (YYYYMMDD format or datetime)
            step: Number of steps backward (default: 1)
            adjust: Whether to adjust if date is exactly on a frequency boundary (default: True)

        Returns:
            Previous date in the class frequency (YYYYMMDD format)

        Examples:
            >>> monthly = Monthly(1)
            >>> monthly.prev('20240115')  # Get previous month-end
            '20231229'

        Raises:
            FrequencyError: If calculation fails or step is invalid
        """
        if step < 1:
            raise FrequencyError(f"Step must be positive, got {step}")

        date = dt_handle(date)
        begin_dt = step_trade_dt(date, -DAILY_SEARCH_STEPS)
        df = self.get(begin_dt, date).tolist()

        try:
            if df and df[-1] == date:
                if adjust and len(df) > step:
                    return df[-step - 1]
                elif len(df) > step:
                    return df[-step - 1]

            if df and len(df) >= step:
                return df[-step]
            elif df:
                return df[0]
            else:
                raise FrequencyError(f"No frequency dates found before {date}")
        except IndexError as e:
            if df:
                return df[0]
            raise FrequencyError(f"Cannot get previous date from {date}: {e}")

    @staticmethod
    def freq_handle(arg, df, step=1):
        if arg == 1:
            tmp_df = df.map(lambda x: adjust_trade_dt(x[:6] + '01', 'next'))
        else:
            tmp_df = df.map(lambda x: step_trade_dt(
                str(int(x[:6]) + step) + '01', -1))

        return tmp_df

    @staticmethod
    def df_handle(begin_dt=DEFAULT_BEGIN_DATE, end_dt=DEFAULT_END_DATE, func=None):
        begin_dt = dt_handle(begin_dt)
        end_dt = dt_handle(end_dt)
        df = get_trade_dts(begin_dt, end_dt).copy()

        df = df.map(func)
        df.drop_duplicates(inplace=True)
        return df

    def _get(self,
             begin_dt=DEFAULT_BEGIN_DATE,
             end_dt=DEFAULT_END_DATE,
             func=None,
             offset=None):
        begin_dt, end_dt = dt_handle(begin_dt), dt_handle(end_dt)
        df = get_trade_dts(
            step_trade_dt(begin_dt, -1 * offset),
            step_trade_dt(end_dt, offset)).to_frame('trade_dt')
        df['_trade_dt'] = pd.to_datetime(df['trade_dt'])
        df['period'] = df['_trade_dt'].map(func)
        all_trade_dt = pd.Series(dtype=object)
        for arg in self.args:
            if arg == 1:
                tmp_df = df.drop_duplicates('period', keep='first')['trade_dt']
            else:
                tmp_df = df.drop_duplicates('period', keep='last')['trade_dt']
            all_trade_dt = pd.concat([all_trade_dt, tmp_df])
        all_trade_dt.sort_values(inplace=True)
        all_trade_dt = all_trade_dt[
            (all_trade_dt >= begin_dt)
            & (all_trade_dt <= end_dt)].drop_duplicates()
        all_trade_dt.reset_index(drop=True, inplace=True)
        return all_trade_dt


class Daily(RefreshBase):
    def get(self, begin_dt=DEFAULT_BEGIN_DATE, end_dt=DEFAULT_END_DATE):
        """[get trade_dt Series with class freq]

        Arguments:
            RefreshBase {[cls]} -- [refreshbase]

        Keyword Arguments:
            begin_dt {str or datetime} -- [begin_dt] (default: {'19900101'})
            end_dt {str or datetime} -- [end_dt] (default: {'20990101'})

        Returns:
            [pd.Series] -- [trade_dt Series]
        """
        return get_trade_dts(begin_dt, end_dt).copy()


class Monthly(RefreshBase):
    def get(self, begin_dt=DEFAULT_BEGIN_DATE, end_dt=DEFAULT_END_DATE):
        """[get trade_dt Series with class freq]

        Arguments:
            RefreshBase {[cls]} -- [refreshbase]

        Keyword Arguments:
            begin_dt {str or datetime} -- [begin_dt] (default: {'19900101'})
            end_dt {str or datetime} -- [end_dt] (default: {'20990101'})

        Returns:
            [pd.Series] -- [trade_dt Series]
        """
        def func(x):
            return f"{x.year}{x.month}"

        return self._get(
            begin_dt=begin_dt, end_dt=end_dt, func=func, offset=MONTHLY_OFFSET)


class Weekly(RefreshBase):
    def get(self, begin_dt=DEFAULT_BEGIN_DATE, end_dt=DEFAULT_END_DATE):
        """[get trade_dt Series with class freq]

        Arguments:
            RefreshBase {[cls]} -- [refreshbase]

        Keyword Arguments:
            begin_dt {str or datetime} -- [begin_dt] (default: {'19900101'})
            end_dt {str or datetime} -- [end_dt] (default: {'20990101'})

        Returns:
            [pd.Series] -- [trade_dt Series]
        """
        def func(x):
            tmpx = x.isocalendar()
            return f"{tmpx[0]}{tmpx[1]}"

        return self._get(
            begin_dt=begin_dt, end_dt=end_dt, func=func, offset=WEEKLY_OFFSET)


class BiWeekly(RefreshBase):
    def get(self, begin_dt=DEFAULT_BEGIN_DATE, end_dt=DEFAULT_END_DATE):
        """[get trade_dt Series with class freq]

        Arguments:
            RefreshBase {[cls]} -- [refreshbase]

        Keyword Arguments:
            begin_dt {str or datetime} -- [begin_dt] (default: {'19900101'})
            end_dt {str or datetime} -- [end_dt] (default: {'20990101'})

        Returns:
            [pd.Series] -- [trade_dt Series]
        """

        all_trade_dt = pd.Series()
        for arg in self.args:
            if arg == 1:
                tmp_df = Weekly(1).get(begin_dt, end_dt)[::2]
            else:
                tmp_df = Weekly(-1).get(begin_dt, end_dt)[::2]
            all_trade_dt = pd.concat([all_trade_dt, tmp_df])
        all_trade_dt.sort_values(inplace=True)
        all_trade_dt.drop_duplicates(inplace=True)
        all_trade_dt.reset_index(drop=True, inplace=True)
        return all_trade_dt


class Quarterly(RefreshBase):
    def get(self, begin_dt=DEFAULT_BEGIN_DATE, end_dt=DEFAULT_END_DATE):
        """[get trade_dt Series with class freq]

        Arguments:
            RefreshBase {[cls]} -- [refreshbase]

        Keyword Arguments:
            begin_dt {str or datetime} -- [begin_dt] (default: {'19900101'})
            end_dt {str or datetime} -- [end_dt] (default: {'20990101'})

        Returns:
            [pd.Series] -- [trade_dt Series]
        """

        def func(x):
            return f"{x.year}{x.quarter}"

        return self._get(
            begin_dt=begin_dt, end_dt=end_dt, func=func, offset=QUARTERLY_OFFSET)


class Reportly(RefreshBase):
    @staticmethod
    def _report(x):
        # Convert to integer for proper date comparison
        date_int = int(x)
        if date_int <= int(x[:4] + '0430'):
            return str(int(x[:4]) - 1) + '11'
        elif date_int <= int(x[:4] + '0831'):
            return x[:4] + '05'
        elif date_int <= int(x[:4] + '1031'):
            return x[:4] + '09'
        elif date_int <= int(x[:4] + '1231'):
            return x[:4] + '11'

    def get(self, begin_dt=DEFAULT_BEGIN_DATE, end_dt=DEFAULT_END_DATE):
        """[get trade_dt Series with class freq]

        Arguments:
            RefreshBase {[cls]} -- [refreshbase]

        Keyword Arguments:
            begin_dt {str or datetime} -- [begin_dt] (default: {'19900101'})
            end_dt {str or datetime} -- [end_dt] (default: {'20990101'})

        Returns:
            [pd.Series] -- [trade_dt Series]
        """

        begin_dt, end_dt = dt_handle(begin_dt), dt_handle(end_dt)
        df = self.df_handle(begin_dt, end_dt, self._report)
        all_trade_dt = pd.Series()
        for arg in self.args:
            if arg == 1:
                tmp_df = df.map(
                    lambda x: adjust_trade_dt(x[:6] + '01', 'next'))
            else:
                def neg_report(x):
                    if x[-2:] == '11':
                        return step_trade_dt(str(int(x[:4]) + 1) + '0501', -1)
                    elif x[-2:] == '09':
                        return step_trade_dt(x[:4] + '1101', -1)
                    elif x[-2:] == '05':
                        return step_trade_dt(x[:4] + '0901', -1)

                tmp_df = df.map(neg_report)
            all_trade_dt = pd.concat([all_trade_dt, tmp_df])
        all_trade_dt.sort_values(inplace=True)
        all_trade_dt = all_trade_dt[(all_trade_dt >= begin_dt)
                                    & (all_trade_dt <= end_dt)].copy()
        all_trade_dt.reset_index(drop=True, inplace=True)
        return all_trade_dt


class Yearly(RefreshBase):
    def get(self, begin_dt=DEFAULT_BEGIN_DATE, end_dt=DEFAULT_END_DATE):
        """[get trade_dt Series with class freq]

        Arguments:
            RefreshBase {[cls]} -- [refreshbase]

        Keyword Arguments:
            begin_dt {str or datetime} -- [begin_dt] (default: {'19900101'})
            end_dt {str or datetime} -- [end_dt] (default: {'20990101'})

        Returns:
            [pd.Series] -- [trade_dt Series]
        """
        def func(x):
            return f"{x.year}"

        return self._get(
            begin_dt=begin_dt, end_dt=end_dt, func=func, offset=YEARLY_OFFSET)


class Halfyearly(RefreshBase):
    @staticmethod
    def _year(x):
        # Convert to integer for proper date comparison
        date_int = int(x)
        if date_int <= int(x[:4] + '0630'):
            return x[:4] + '01'
        elif date_int <= int(x[:4] + '1231'):
            return x[:4] + '07'

    def get(self, begin_dt=DEFAULT_BEGIN_DATE, end_dt=DEFAULT_END_DATE):
        """[get trade_dt Series with class freq]

        Arguments:
            RefreshBase {[cls]} -- [refreshbase]

        Keyword Arguments:
            begin_dt {str or datetime} -- [begin_dt] (default: {'19900101'})
            end_dt {str or datetime} -- [end_dt] (default: {'20990101'})

        Returns:
            [pd.Series] -- [trade_dt Series]
        """

        begin_dt, end_dt = dt_handle(begin_dt), dt_handle(end_dt)
        df = self.df_handle(begin_dt, end_dt, self._year)
        all_trade_dt = pd.Series()
        for arg in self.args:
            tmp_df = self.freq_handle(arg, df, 6)
            all_trade_dt = pd.concat([all_trade_dt, tmp_df])
        all_trade_dt.sort_values(inplace=True)
        all_trade_dt = all_trade_dt[(all_trade_dt >= begin_dt)
                                    & (all_trade_dt <= end_dt)].copy()
        all_trade_dt.reset_index(drop=True, inplace=True)
        return all_trade_dt


if __name__ == '__main__':
    pass